import burp.api.montoya.http.message.HttpHeader;
import burp.api.montoya.http.message.HttpRequestResponse;
import burp.api.montoya.http.message.requests.HttpRequest;
import burp.api.montoya.ui.contextmenu.ContextMenuEvent;
import burp.api.montoya.ui.contextmenu.ContextMenuItemsProvider;
import ui.PermissionTestPanel;
import utils.LoggerUtil;

import javax.swing.*;
import java.awt.*;
import java.util.ArrayList;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;

@SuppressWarnings("ClassCanBeRecord")
public class MenuItemProvider implements ContextMenuItemsProvider {

    private final PermissionTestPanel permissionTestPanel;

    private enum MenuItemType {
        REQUEST,
        COOKIE
    }

    public MenuItemProvider(PermissionTestPanel permissionTestPanel) {
        this.permissionTestPanel = permissionTestPanel;
    }

    @Override
    public List<Component> provideMenuItems(ContextMenuEvent event) {
        // 创建菜单项
        List<Component> menuItems = new ArrayList<>();

        // 添加按钮"发送请求到插件面板"
        JMenuItem requestToPermissionTest = createMenuItem("Send Request to PermissionTest", event, MenuItemType.REQUEST);

        // 添加按钮"发送cookie到插件面板"
        JMenuItem cookieToPermissionTest = createMenuItem("Send Cookie to PermissionTest", event, MenuItemType.COOKIE);

        // 添加按钮到菜单项列表
        menuItems.add(requestToPermissionTest);
        menuItems.add(cookieToPermissionTest);
        return menuItems;
    }

    // 创建菜单项
    private JMenuItem createMenuItem(String text, ContextMenuEvent event, MenuItemType type) {
        JMenuItem menuItem = new JMenuItem(text);
        menuItem.addActionListener(e -> {
            List<HttpRequestResponse> selectedRequests = new ArrayList<>();

            // 从Proxy History获取请求和响应
            if (event.selectedRequestResponses() != null) {
                selectedRequests.addAll(event.selectedRequestResponses());
            }

            // 从Repeater或Intruder获取请求
            if (selectedRequests.isEmpty() && event.messageEditorRequestResponse().isPresent()) {
                HttpRequest currentRequest = event.messageEditorRequestResponse().get().requestResponse().request();
                if (currentRequest != null) {
                    HttpRequestResponse request = HttpRequestResponse.httpRequestResponse(
                            currentRequest, null);
                    selectedRequests.add(request);
                }
            }

            if (!selectedRequests.isEmpty()) {
                switch (type) {
                    case REQUEST:
                        permissionTestPanel.addRequestsToTable(selectedRequests);
                        LoggerUtil.logOutput("Selected " + selectedRequests.size() + " requests.");
                        break;
                    case COOKIE:
                        // 清空临时请求头
                        permissionTestPanel.clearTempHeaderContent();
                        // 使用LinkedHashMap保持插入顺序并自动去重
                        Map<String, String> cookieMap = new LinkedHashMap<>();

                        // 遍历所有选中的请求
                        for (HttpRequestResponse requestResponse : selectedRequests) {
                            HttpRequest request = requestResponse.request();
                            if (request != null) {
                                // 获取所有的Cookie头
                                List<String> cookieHeaders = request.headers().stream()
                                        .filter(header -> "Cookie".equalsIgnoreCase(header.name()))
                                        .map(HttpHeader::value)
                                        .collect(ArrayList::new, ArrayList::add, ArrayList::addAll);

                                // 解析每个Cookie头中的键值对
                                for (String cookieHeader : cookieHeaders) {
                                    if (cookieHeader != null && !cookieHeader.isEmpty()) {
                                        // 按分号分割Cookie键值对
                                        String[] cookiePairs = cookieHeader.split(";");
                                        for (String cookiePair : cookiePairs) {
                                            String trimmedCookiePair = cookiePair.trim();
                                            if (!trimmedCookiePair.isEmpty()) {
                                                // 分割键和值
                                                int equalsIndex = trimmedCookiePair.indexOf('=');
                                                if (equalsIndex > 0) {
                                                    String cookieName = trimmedCookiePair.substring(0, equalsIndex).trim();
                                                    String cookieValue = trimmedCookiePair.substring(equalsIndex + 1).trim();
                                                    // 用新值覆盖旧值
                                                    cookieMap.put(cookieName, cookieValue);
                                                } else {
                                                    // 没有值的Cookie
                                                    cookieMap.put(trimmedCookiePair, "");
                                                }
                                            }
                                        }
                                    }
                                }
                            }
                        }

                        // 构建最终的Cookie头
                        if (!cookieMap.isEmpty()) {
                            StringBuilder content = new StringBuilder();
                            content.append("Cookie: ");
                            boolean first = true;
                            for (Map.Entry<String, String> entry : cookieMap.entrySet()) {
                                if (!first) {
                                    content.append("; ");
                                }
                                if (entry.getValue().isEmpty()) {
                                    content.append(entry.getKey());
                                } else {
                                    content.append(entry.getKey()).append("=").append(entry.getValue());
                                }
                                first = false;
                            }
                            permissionTestPanel.setTempHeaderContent(content.toString());
                        }
                        LoggerUtil.logOutput("Sent " + selectedRequests.size() + " cookies to PermissionTest panel.");
                        break;

                }
            } else {
                LoggerUtil.logOutput("No requests selected or available in current context.");
            }
        });
        return menuItem;
    }
}
